package me.coley.recaf.parse.bytecode;

import javassist.ClassPool;
import javassist.CtBehavior;
import javassist.CtClass;
import me.coley.recaf.compiler.JavassistASMTranslator;
import me.coley.recaf.compiler.JavassistCompilationResult;
import me.coley.recaf.compiler.JavassistCompiler;
import me.coley.recaf.control.Controller;
import me.coley.recaf.metadata.Comments;
import me.coley.recaf.parse.bytecode.ast.AST;
import me.coley.recaf.parse.bytecode.ast.LabelAST;
import me.coley.recaf.parse.bytecode.ast.MethodDefinitionAST;
import me.coley.recaf.parse.bytecode.ast.RootAST;
import me.coley.recaf.parse.bytecode.exception.AssemblerException;
import me.coley.recaf.util.TypeUtil;
import org.objectweb.asm.tree.AbstractInsnNode;
import org.objectweb.asm.tree.LabelNode;
import org.objectweb.asm.tree.LocalVariableNode;
import org.objectweb.asm.tree.MethodNode;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.util.*;

/**
 * AST compilation context.
 *
 * @author xxDark
 */
public final class MethodCompilation {
    private final ParseResult<RootAST> ast;
    private final MethodDefinitionAST methodDefinition;
    private final MethodNode node;
    private final String declaringType;
    private final Controller controller;
	private final List<LocalVariableNode> priorVars;
	private final Comments comments = new Comments();
    private final Map<String, LabelNode> nameToLabel = new LinkedHashMap<>();
    private final Map<LabelNode, LabelAST> labelToAst = new HashMap<>();
    private final Map<AbstractInsnNode, AST> insnToAst = new HashMap<>();
    private final Map<Integer, AbstractInsnNode> lineToInsn = new HashMap<>();
	private VariableNameCache variableNames;

	/**
	 * @param ast
	 * 		Root AST.
	 * @param methodDefinition
	 * 		Definition of the method to compile.
	 * @param node
	 * 		ASM representation of the method.
	 * @param declaringType
	 * 		The type declaring the method.
	 * @param controller
	 * 		Controller to pull references from.
	 * @param priorVars
	 * 		Prior variable information.
	 */
	public MethodCompilation(ParseResult<RootAST> ast, MethodDefinitionAST methodDefinition,
							 MethodNode node, String declaringType, Controller controller,
							 List<LocalVariableNode> priorVars) {
		this.ast = ast;
		this.methodDefinition = methodDefinition;
		this.node = node;
		this.declaringType = declaringType;
		this.controller = controller;
		this.priorVars = priorVars;
	}

    /**
     * @return  root AST parse result.
     */
    public ParseResult<RootAST> getAst() {
        return ast;
    }

    /**
     * @return method definition.
     */
    public MethodDefinitionAST getMethodDefinition() {
        return methodDefinition;
    }

    /**
     * @return ASM representation of the method.
     */
    public MethodNode getNode() {
        return node;
    }

    /**
     * @param name
     *      Name of the label.
     *
     * @return label by it's name. May be {@code null}.
     */
    public LabelNode getLabel(String name) {
        return nameToLabel.get(name);
    }

    /**
     * @param node
     *      Label node.
     *
     * @return label's AST. May be {@code null}.
     */
    public LabelAST getLabelAst(LabelNode node) {
        return labelToAst.get(node);
    }

    /**
     * @return a map of name > label layout.
     */
    public Map<String, LabelNode> getNameToLabel() {
        return Collections.unmodifiableMap(nameToLabel);
    }

    /**
     * @return a map of label > AST layout.
     */
    public Map<LabelNode, LabelAST> getLabelToAst() {
        return Collections.unmodifiableMap(labelToAst);
    }

    /**
     * Adds and register instruction to specific AST.
     *
     * @param insn
     *      Instruction to assign.
     * @param ast
     *      Instruction's AST.
     *
     * @see MethodCompilation#assignInstruction(AbstractInsnNode, AST)
     */
    public void addInstruction(AbstractInsnNode insn, AST ast) {
        node.instructions.add(insn);
        assignInstruction(insn, ast);
    }

    /**
     * Adds a comment at the current instruction offset.
     *
     * @param comment
     * 		Comment string to add.
     */
    public void addComment(String comment) {
        int index = node.instructions.size();
		comments.addComment(index, comment);
	}

	/**
	 * Compiles an expression and inserts the generated code at the current instruction offset.
	 *
	 * @param expression
	 * 		Expression to compile.
	 * @param ast
	 * 		Associated AST.
	 *
	 * @throws AssemblerException
	 * 		When the expression cannot be compiled.
	 */
	public void addExpression(String expression, AST ast) throws AssemblerException {
		// TODO: Since instructions and AST have a one to one mapping we may want to make a map of lines
		//       to expressions if something goes wrong in the code generated by the expression.
		try {
			InputStream stream = new ByteArrayInputStream(controller.getWorkspace().getRawClass(declaringType));
			CtClass declaring = ClassPool.getDefault().makeClass(stream);
			CtBehavior containerMethod;
			if (node.name.equals("<init>")) {
				containerMethod = declaring.getConstructor(node.desc);
			} else {
				containerMethod = declaring.getMethod(node.name, node.desc);
			}
			// Compile with Javassist
			JavassistCompilationResult result =
					JavassistCompiler.compileExpression(declaring, containerMethod, expression,
							priorVars, variableNames);
			// Translate to ASM
			JavassistASMTranslator translator = new JavassistASMTranslator();
			translator.visit(declaring, result.getBytecode().toCodeAttribute());
			node.instructions.add(translator.getInstructions());
			// Add variables
			Object[] entries = nameToLabel.keySet().toArray();
			result.getSymbols().forEach((name, dec) -> {
				String desc = dec.getClassName();
				if (TypeUtil.isPrimitiveClassName(desc))
					desc = TypeUtil.classToPrimitive(desc);
				int index = dec.getLocalVar();
				LabelNode start = nameToLabel.get(String.valueOf(entries[0]));
				LabelNode end = nameToLabel.get(String.valueOf(entries[nameToLabel.size() - 1]));
				node.localVariables.add(new LocalVariableNode(name, desc, null, start, end, index));
			});
		} catch (Exception ex) {
			String msg = ex.getMessage();
			if (msg == null)
				msg = "Unknown error with expression compilation";
			throw new AssemblerException(ex, msg, ast.getLine());
		}
	}

	/**
	 * Assigns instruction to specific AST.
	 * Also sets the line to specific instruction.
	 *
	 * @param insn
	 * 		Instruction to assign.
	 * @param ast
	 * 		Instruction's AST.
	 */
	public void assignInstruction(AbstractInsnNode insn, AST ast) {
		insnToAst.put(insn, ast);
		lineToInsn.put(ast.getLine(), insn);
	}

    /**
     * Assigns label to specific AST.
     *
     * @param label
     *      Instruction to assign.
     * @param ast
     *      Label's AST.
     */
    public void assignLabel(LabelNode label, LabelAST ast) {
        nameToLabel.put(ast.getName().getName(), label);
        labelToAst.put(label, ast);
    }

    /**
     * @param node
     *      Instruction to get line from.
     *
     * @return line of the instruction. May return {@code -1} if AST is missing.
     */
    public int getLine(AbstractInsnNode node) {
        AST ast = insnToAst.get(node);
        return ast == null ? -1 : ast.getLine();
    }

    /**
     * @param line
     *      Line to get instruction from.
     * @return instruction by it's line. May be {@code null}.
     */
    public AbstractInsnNode getInsn(int line) {
        return lineToInsn.get(line);
    }

    /**
     * @return Name cache of variables for the method.
     */
    public VariableNameCache getVariableNameCache() {
        return variableNames;
    }

    // Internal methods
    void setVariableNames(VariableNameCache variableNames) {
        this.variableNames = variableNames;
    }

    void onCompletion() {
        comments.applyTo(node);
    }
}
